package com.hotshow.caleb.drivedetectiondemo;

import static com.hotshow.caleb.drivedetectiondemo.commonconst.DetectionConst.MINIMUM_CONFIDENCE_TF_OD_API;
import static com.hotshow.caleb.drivedetectiondemo.commonconst.DetectionConst.TEXT_SIZE_DIP;

import android.content.pm.ActivityInfo;
import android.content.res.Configuration;
import android.graphics.Bitmap;
import android.graphics.Canvas;
import android.graphics.Color;
import android.graphics.Matrix;
import android.graphics.Paint;
import android.graphics.RectF;
import android.graphics.Typeface;
import android.graphics.drawable.Drawable;
import android.os.Bundle;
import android.os.Handler;
import android.os.HandlerThread;
import android.util.DisplayMetrics;
import android.util.Log;
import android.util.TypedValue;
import android.view.Surface;
import android.view.SurfaceHolder;
import android.view.WindowManager;
import android.widget.Toast;

import androidx.annotation.NonNull;
import androidx.annotation.Nullable;
import androidx.appcompat.app.AppCompatActivity;

import com.bumptech.glide.Glide;
import com.bumptech.glide.load.engine.DiskCacheStrategy;
import com.bumptech.glide.request.target.CustomTarget;
import com.bumptech.glide.request.transition.Transition;
import com.hotshow.caleb.drivedetectiondemo.commonconst.DetectionConst;
import com.hotshow.caleb.drivedetectiondemo.customeView.OverlayView;
import com.hotshow.caleb.drivedetectiondemo.databinding.ActivityMainBinding;
import com.hotshow.caleb.drivedetectiondemo.env.BorderedText;
import com.hotshow.caleb.drivedetectiondemo.tflite.Classifier;
import com.hotshow.caleb.drivedetectiondemo.tflite.DetectorFactory;
import com.hotshow.caleb.drivedetectiondemo.tflite.YoloV5Classifier;
import com.hotshow.caleb.drivedetectiondemo.tracking.MultiBoxTracker;
import com.hotshow.caleb.drivedetectiondemo.utils.ImageUtils;
import com.hotshow.caleb.drivedetectiondemo.utils.Logger;
import com.hotshow.caleb.drivedetectiondemo.utils.Utils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

public class MainActivity extends AppCompatActivity implements SurfaceHolder.Callback {
    private ActivityMainBinding binding;
    private static final Logger LOGGER = new Logger();
    private static final String ASSET_PATH = "";
    protected ArrayList<String> modelStrings = new ArrayList<>();
    private YoloV5Classifier detector;
    //源图片
    private Bitmap rgbFrameBitmap = null;
    //变换图片
    private Bitmap croppedBitmap = null;
    //最终图片
    private Bitmap cropCopyBitmap = null;
    //帧到裁剪的变换矩阵
    private Matrix frameToCropTransform;
    //裁剪到帧的变换矩阵
    private Matrix cropToFrameTransform;
    //裁剪图片时是否保持纵横比
    private static final boolean MAINTAIN_ASPECT = true;
    OverlayView trackingOverlay;
    private boolean computingDetection = false;
    private MultiBoxTracker tracker;
    private int imageWidth;
    private int imageHeight;
    private static final ScheduledExecutorService sc = Executors.newScheduledThreadPool(4);
    public static final ExecutorService es = Executors.newFixedThreadPool(4);
    private SurfaceHolder surfaceHolder;
    private BorderedText borderedText;

    private Handler handler;
    private HandlerThread handlerThread;

    @Override
    protected void onCreate(Bundle savedInstanceState) {
        super.onCreate(savedInstanceState);
        getWindow().addFlags(WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON);
        binding = ActivityMainBinding.inflate(getLayoutInflater());
        setContentView(binding.getRoot());
        Utils.makeSystemBarsTransparent(this);
        forceScreenOrientation();
        initSurfaceView();
        modelStrings = Utils.getModelStrings(getAssets(), ASSET_PATH);
        CompletableFuture<Bitmap> pictureDisplay = getPictureDisplayParams();
        pictureDisplay.thenAccept(result -> {
            imageWidth = result.getWidth();
            imageHeight = result.getHeight();
            initAIModel(4);
            initTracker();
        });
    }


    private void initSurfaceView() {
        surfaceHolder = binding.surface.getHolder();
        surfaceHolder.addCallback(this);
    }


    /**
     * 获取图片的宽和高
     */
    private CompletableFuture<Bitmap> getPictureDisplayParams() {
        final CompletableFuture<Bitmap> future = new CompletableFuture();
        Glide.with(MainActivity.this).asBitmap().skipMemoryCache(true).diskCacheStrategy(DiskCacheStrategy.NONE).load(DetectionConst.CAMERA_URL)
                .into(new CustomTarget<Bitmap>() {
                    @Override
                    public void onResourceReady(@NonNull Bitmap bitmap, @Nullable Transition<? super Bitmap> transition) {
                        future.complete(bitmap);
                    }

                    @Override
                    public void onLoadCleared(@Nullable Drawable drawable) {
                        Log.d("dede", "onLoadCleared: ");
                        future.completeExceptionally(new Exception("Load Failure"));
                    }
                });
        return future;
    }

    private void initAIModel(int numThreads) {
        handler.post(() -> {
            // Disable classifier while updating
            if (detector != null) {
                detector.close();
                detector = null;
            }
            //加载模型
            try {
                detector = DetectorFactory.getDetector(getAssets(), !modelStrings.isEmpty() ? modelStrings.get(0) : "");
                if (detector == null) {
                    return;
                }
            } catch (IOException e) {
                e.printStackTrace();
                Toast toast = Toast.makeText(getApplicationContext(), "Classifier could not be initialized", Toast.LENGTH_SHORT);
                toast.show();
                finish();
            }
            //使用CPU和NNAP完成推理，因为模拟器不含有gpu,所以不能使用gpu,但是gpu的推理速度远远优于多核心cpu
            detector.useCPU();
            detector.setNumThreads(numThreads);
            detector.useNNAPI();
            int cropSize = detector.getInputSize();
            croppedBitmap = Bitmap.createBitmap(cropSize, cropSize, Bitmap.Config.ARGB_8888);
            frameToCropTransform =
                    ImageUtils.getTransformationMatrix(
                            imageWidth, imageHeight,
                            cropSize, cropSize,
                            0, MAINTAIN_ASPECT);
            cropToFrameTransform = new Matrix();
            frameToCropTransform.invert(cropToFrameTransform);
        });

    }


    public int[] getRgbBytesFromBitmap(Bitmap bitmap) {
          int width = bitmap.getWidth();
        int height = bitmap.getHeight();
        int[] rgbBytes = new int[width * height];
        // 直接获取像素数据到 int 数组
        bitmap.getPixels(rgbBytes, 0, width, 0, 0, width, height);
        return rgbBytes;
    }


    @Override
    public  void surfaceCreated(@NonNull SurfaceHolder holder) {
        sc.scheduleAtFixedRate(() -> {
            CompletableFuture<Bitmap> pictureDisplayParams = getPictureDisplayParams();
            CompletableFuture.anyOf(pictureDisplayParams).thenRun(() -> {
                try {
                    Bitmap bitmap = pictureDisplayParams.get();
                    drawBitmap(bitmap);
                    //把图片传给模型去处理
                    processImage(bitmap);
                } catch (ExecutionException e) {
                    throw new RuntimeException(e);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            });
        }, 0, 30, TimeUnit.MILLISECONDS);
    }


    private synchronized void drawBitmap(Bitmap bitmap) {
        Canvas canvas = null;
        try {
            canvas = surfaceHolder.lockCanvas();
            if (canvas != null) {
                synchronized (surfaceHolder) {
                    canvas.drawBitmap(bitmap, 0, 0, null);
                }
            }
        } finally {
            if (canvas != null) {
                surfaceHolder.unlockCanvasAndPost(canvas);
            }
        }
    }


    @Override
    public void surfaceChanged(@NonNull SurfaceHolder holder, int format, int width, int height) {

    }

    @Override
    public void surfaceDestroyed(@NonNull SurfaceHolder holder) {

    }


    private void processImage(Bitmap bitmap) {
        trackingOverlay.postInvalidate();
        rgbFrameBitmap.setPixels(getRgbBytesFromBitmap(bitmap), 0, imageWidth, 0, 0, imageWidth, imageHeight);
        Canvas canvas1 = new Canvas(croppedBitmap);
        canvas1.drawBitmap(rgbFrameBitmap, frameToCropTransform, null);
        runInBackground(new Runnable() {
            @Override
            public void run() {
                final List<Classifier.Recognition> results = detector.recognizeImage(croppedBitmap);
                cropCopyBitmap = Bitmap.createBitmap(croppedBitmap);
                final Canvas canvas = new Canvas(cropCopyBitmap);
                final Paint paint = new Paint();
                paint.setColor(Color.RED);
                paint.setStyle(Paint.Style.STROKE);
                paint.setStrokeWidth(2.0f);
                float minimumConfidence = MINIMUM_CONFIDENCE_TF_OD_API;
                final List<Classifier.Recognition> mappedRecognitions =
                        new LinkedList<Classifier.Recognition>();
                for (final Classifier.Recognition result : results) {
                    final RectF location = result.getLocation();
                    if (location != null && result.getConfidence() >= minimumConfidence) {
                        Log.d("dede", "run: "+result);
                        canvas.drawRect(location, paint);
                        cropToFrameTransform.mapRect(location);
                        result.setLocation(location);
                        mappedRecognitions.add(result);
                    }
                }
                tracker.processResults(mappedRecognitions);
                trackingOverlay.postInvalidate();
            }
        });
    }

    private void initTracker() {
        final float textSizePx =
                TypedValue.applyDimension(
                        TypedValue.COMPLEX_UNIT_DIP, TEXT_SIZE_DIP, getResources().getDisplayMetrics());
        borderedText = new BorderedText(textSizePx);
        borderedText.setTypeface(Typeface.MONOSPACE);
        tracker = new MultiBoxTracker(this);
        final String modelString = modelStrings.get(0);

        try {
            detector = DetectorFactory.getDetector(getAssets(), modelString);
        } catch (final IOException e) {
            e.printStackTrace();
            LOGGER.e(e, "Exception initializing classifier!");
            Toast toast =
                    Toast.makeText(
                            getApplicationContext(), "Classifier could not be initialized", Toast.LENGTH_SHORT);
            toast.show();
            finish();
        }
        int cropSize = detector.getInputSize();
        //相机原始的图片大小
        rgbFrameBitmap = Bitmap.createBitmap(imageWidth, imageHeight, Bitmap.Config.ARGB_8888);
        //需要传给AI的图片大小
        croppedBitmap = Bitmap.createBitmap(cropSize, cropSize, Bitmap.Config.ARGB_8888);
        //todo 此处的传感器角度直接是写死的90
        //todo 保持纵横比地缩放图片
        frameToCropTransform =
                ImageUtils.getTransformationMatrix(
                        imageWidth, imageHeight,
                        cropSize, cropSize,
                        getScreenOrientation(), MAINTAIN_ASPECT);

        cropToFrameTransform = new Matrix();
        frameToCropTransform.invert(cropToFrameTransform);
        trackingOverlay = (OverlayView) findViewById(R.id.tracking_overlay);
        trackingOverlay.addCallback(
                new OverlayView.DrawCallback() {
                    @Override
                    public void drawCallback(final Canvas canvas) {
                        tracker.draw(canvas);
                    }
                });

        tracker.setFrameConfiguration(imageWidth, imageHeight, getScreenOrientation());
    }

    //获取屏幕的旋转角度
    protected int getScreenOrientation() {
        switch (getWindowManager().getDefaultDisplay().getRotation()) {
            case Surface.ROTATION_270:
                return 270;
            case Surface.ROTATION_180:
                return 180;
            case Surface.ROTATION_90:
                return 90;
            default:
                return 0;
        }
    }


    private void forceScreenOrientation() {
        DisplayMetrics displayMetrics = new DisplayMetrics();
        getWindowManager().getDefaultDisplay().getMetrics(displayMetrics);
        int width = displayMetrics.widthPixels;
        int height = displayMetrics.heightPixels;
        int orientation = getResources().getConfiguration().orientation;
        // 判断设备类型（手机或平板）
        // 需要在 values/bools.xml 中定义 isTablet
        if (Utils.isTablet(this)) {
            // 平板：默认纵向，除非当前方向为横向
            if (orientation == Configuration.ORIENTATION_LANDSCAPE) {
                setRequestedOrientation(ActivityInfo.SCREEN_ORIENTATION_LANDSCAPE);
            } else {
                setRequestedOrientation(ActivityInfo.SCREEN_ORIENTATION_PORTRAIT);
            }
        } else {
            // 手机：强制横屏
            if (width > height) {
                setRequestedOrientation(ActivityInfo.SCREEN_ORIENTATION_LANDSCAPE);
            } else {
                setRequestedOrientation(ActivityInfo.SCREEN_ORIENTATION_REVERSE_LANDSCAPE); // 反向横屏
            }
        }
    }

    @Override
    public synchronized void onResume() {
        super.onResume();
        handlerThread = new HandlerThread("inference");
        handlerThread.start();
        handler = new Handler(handlerThread.getLooper());
    }

    @Override
    public synchronized void onPause() {
        LOGGER.d("onPause " + this);
        handlerThread.quitSafely();
        try {
            handlerThread.join();
            handlerThread = null;
            handler = null;
        } catch (final InterruptedException e) {
            LOGGER.e(e, "Exception!");
        }
        super.onPause();
    }


    protected synchronized void runInBackground(final Runnable r) {
        if (handler != null) {
            handler.post(r);
        }
    }


    @Override
    public void onConfigurationChanged(@NonNull Configuration newConfig) {
        super.onConfigurationChanged(newConfig);
        forceScreenOrientation();
    }
}